import pandas as pd
from ted_funs import *
import matplotlib.pyplot as plt
plt.rcParams['font.family'] =  ['SimHei']
plt.rcParams['axes.unicode_minus'] = False 

# Load the Excel file
file_path = '/Users/xx/Desktop/股票/股票1230.xlsx'
df = pd.read_excel(file_path, sheet_name='ALL')
df_name = df


# Filter rows where flag is 1 and select the 'ts_code' column
lst = df[df['Flag'] == 1]['ts_code'].tolist()

df_lst = []

for ts_code in lst:
	print(ts_code)
	query = "select * from daily_2024 where ts_code = '{}' and trade_date LIKE '2023-12-%';".format(ts_code)
	df = mysql_read(query)
	df_lst.append(df)

combined_df = pd.concat(df_lst)


df = combined_df
df['trade_date'] = df['trade_date'].astype(str)
df['trade_date'] = df['trade_date'].str[-2:]

unique_ts_codes = df['ts_code'].unique()
num_ts_codes = len(unique_ts_codes)
total = num_ts_codes + 1

cols = 3  # You can adjust this as needed
rows = total // cols + (total % cols > 0)



fig = plt.figure(figsize=(25, 5 * rows))  # Adjust figure size as needed

ax1 = fig.add_subplot(rows, cols, 1)

for ts_code in unique_ts_codes:
    name = df_name[df_name['ts_code'] == ts_code]['name'].values[0]

    subset = df[df['ts_code'] == ts_code]
    ax1.plot(subset['trade_date'], subset['pct_chg'], label=name)

plt.xlabel('日期')
plt.ylabel('涨跌幅')
plt.title('股票集')
plt.legend()



for index, ts_code in enumerate(unique_ts_codes, start=2):
    ax = fig.add_subplot(rows, cols, index)

    subset = df[df['ts_code'] == ts_code]

    ax.plot(subset['trade_date'], subset['pct_chg'], linestyle="dashed", marker="o")
    ax.axhline(0, color='red', linestyle='--', label='Horizontal Line at 0')

  
    plt.xlabel('日期')
    plt.ylabel('涨跌幅')

    name = df_name[df_name['ts_code'] == ts_code]['name'].values[0]
    plt.title(f'{name}')


plt.tight_layout()  # Adjusts the plots so they don't overlap

fig.savefig('sine_wave.png')
plt.show()
print('finished...')

